// Copyright(c) 2020-2025 Hongjing
// All rights reserved.
//
// Author: Yang Zhu
// Update: 2020-04-07

#ifndef LATTICE_NET_DECODE_H_
#define LATTICE_NET_DECODE_H_

#include <math.h>

#include <Eigen/Dense>
#include <array>
#include <memory>
#include <mutex>
#include <unordered_map>

#include "common/perception_defs.h"
#include "common/infer_common.h"

class LatticeNetDecode {
 public:
  /// @brief max decode result size
  static constexpr std::uint16_t kMaxResultSize{5120U};

  /// @brief
  static constexpr std::uint32_t kMaxObjInfoSize{16};

  /// @brief max lanemarker result size
  static constexpr std::uint32_t kMaxLaneMarkerResultSize{51200U};

  /// @brief
  static constexpr std::uint32_t kMaxMapFeatureSize{6};

  explicit LatticeNetDecode() noexcept;
  ~LatticeNetDecode();
  LatticeNetDecode(LatticeNetDecode const&) = delete;
  LatticeNetDecode& operator=(LatticeNetDecode const&) = delete;
  LatticeNetDecode(LatticeNetDecode&&) = delete;
  LatticeNetDecode& operator=(LatticeNetDecode&&) = delete;

  struct ObjClassInfo {
    std::string name;
    float threshold;
    float anchor_height;
    float length;
    float width;
    float height;
  };

  struct ObjDecodeParameter {
    int32_t bbox_info_number;
    int32_t down_size;
    std::vector<float> net_range;
    std::vector<float> voxel_size;
    std::vector<ObjClassInfo> class_info;
    float nms_iou_threshold;
  };

  struct MapDecodeParameter {
    int32_t map_info_number;
    int32_t down_size;
    std::vector<float> net_range;
    std::vector<float> voxel_size;
    std::vector<float> class_threshold;
    float distance_threshold;
  };

  /// @brief decode object
  /// @param dataRaw infer output raw data
  /// @param parameters decode parameters
  /// @param boundingBoxes bounding box
  /// @param err error message
  /// @return if no error return true, otherwise false
  bool DecodeObject(InferOutputDataRaw const& dataRaw,
                    ObjDecodeParameter const& parameters,
                    FixedArray<BoundingBox, kMaxBoundingBoxNum>& boundingBoxes,
                    std::string& err) noexcept;

  /// @brief decode map
  /// @param dataRaw infer output raw data
  /// @param parameters decode parameters
  /// @param map map
  /// @param err error message
  /// @return if no error return true, otherwise false
  bool DecodeMap(InferOutputDataRaw const& dataRaw,
                 MapDecodeParameter const& parameters, Map& map,
                 std::string& err) noexcept;

 private:
  /// @brief
  /// @param box_raw
  /// @param classBase
  /// @param boundingBoxes
  /// @param err
  /// @return
  bool decode_box(InferBoxRaw const& box_raw, int const& class_number,
                  int const& class_base, ObjDecodeParameter const& parameters,
                  FixedArray<BoundingBox, kMaxBoundingBoxNum>& bounding_boxes,
                  std::string& err) noexcept;

  /// @brief gen polygon
  /// @param bbox bounding box
  /// @param polygon polygon
  void gen_polygon(BoundingBox const& bbox,
                   FixedArray<Point2f, 4>& polygon) noexcept;

 private:
  /// @brief object scores
  std::shared_ptr<std::array<float, kMaxResultSize>> obj_scores_ptr;
  /// @brief object labels
  std::shared_ptr<std::array<int32_t, kMaxResultSize>> obj_labels_ptr;
  /// @brief object boxes
  std::shared_ptr<std::array<float, kMaxResultSize * kMaxObjInfoSize>>
      obj_boxes_ptr;
  /// @brief nms keep
  std::shared_ptr<std::array<bool, kMaxResultSize>> obj_nms_keep_ptr;

  ///@brief map points
  std::shared_ptr<std::array<float, kMaxResultSize * kMaxMapFeatureSize>>
      map_pts_ptr_;
};

#endif  // LATTICE_NET_DECODE_H_
